from contextvars import ContextVar
from time import perf_counter
from starlette.requests import Request
import typing
from starlette.datastructures import Headers
from starlette.types import ASGIApp, Message, Receive, Scope, Send
import shortuuid
from starlette.responses import Response

from .contextvar import log_request_var, logrequest
from .enums import RecordModel
from afast_core.core_plugins.loguru import logger

import json
from typing import Optional
from pydantic import BaseModel, Field


class ResponseInfo(BaseModel):
    headers: Optional[Headers] = Field(default=None, title="Response header")
    body: str = Field(default="", title="Response body")
    status_code: Optional[int] = Field(default=None, title="Status code")

    class Config:
        arbitrary_types_allowed = True


# 存贮日志内容的上下文信息
log_msg_var: ContextVar[dict] = ContextVar("log_msg_var", default={})


class LoguruPluginClientMiddleware:
    '''
    # 这个中间件的话，如果需要在内部-0---读取请求体内容，需要再最终添加，也就是需要在最前面的执行
    # 也就是需要在所有的中间件的最后面再去注册，不能放在其他的前面,
    # 如果需要在内部消费   body = await self.request.body()，则需要开启is_proxy=True
    '''

    def __init__(self, *, app: ASGIApp, is_proxy=True, client) -> None:
        self.app = app
        self.is_proxy = is_proxy
        self.client = client
        self.request: typing.Optional[Request] = None

    async def get_body(self):
        """获取请求BODY，实现使用代理方式解析读，解决在中间件中火球Body的问题"""
        body = await self.request.body()
        return body

    async def get_json(self):
        """获取json请求参数"""
        return json.loads(await self.get_body())

    async def before_request(self, request: Request) -> [Response, None]:
        """如果需要修改请求信息，可直接重写此方法"""
        pass
        # # 追踪ID
        request.state.traceid = str(shortuuid.uuid())
        request.state.traceindex = 0
        #
        request.state.close_record = False
        # # 计算时间
        request.state.start_time = perf_counter()

    async def after_request(self, request: Request, token=None, response: Response = None) -> [Response,
                                                                                               None]:
        """请求后的处理【记录请求耗时等，注意这里没办法对响应结果进行处理】"""
        pass
        # 记录响应报文体内容信息，
        log_msg = log_msg_var.get()
        if self.client.settings.IS_RECORD_RESPONSE and log_msg and response.status_code != 404:
            logger.info(str(response.body, 'utf-8'), event_name='response')

        try:
            request.state.traceindex = None
            request.state.traceid = None
            request.state.trace_logs_record = None
            log_request_var.reset(token)
        except:
            pass

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            return await self.app(scope, receive, send)
        # 和上面的判断类似
        if scope["type"] not in ["http", 'websocket']:
            await self.app(scope, receive, send)
            return

        # 解决读取BODY问题
        if self.is_proxy:
            receive_ = await receive()

            async def receive():
                return receive_

        # 这里考虑直接读取一次body然后保存到对应的上下文中

        # 解析当前的请求体
        self.request = Request(scope, receive=receive)
        # 获取到客户端对象需要过滤的不记录的URL信息，这里直接的跳过
        if self.client.filter_request_url(request=self.request):
            # 新增过滤直接跳过不需要做其他判断处理
            return await self.app(scope, receive, send)

        # 解析报文体内容
        response_info = ResponseInfo()
        # 自定义回调函数，可以自己进行重写实现具体的业务逻辑
        await self.before_request(self.request) or self.app
        # 解析当前的请求体
        token = log_request_var.set(self.request)
        #
        # 离散是日志记录模式
        if self.client.settings.MODEL == RecordModel.SCATTERED:
            logrequest.state.record_model = RecordModel.SCATTERED
            log_msg = await self.client.make_request_log_msg(self.request)
            log_msg_var.set(log_msg)
            # 如果过滤了，则也记录请求信息了
            if log_msg:
                logger.info(log_msg, event_name='request')
        else:
            # 集中式日志记录模式已
            pass
            # 创建全局是日志上下文
            logrequest.state.record_model = RecordModel.CENTRALIZED
            self.request.state.trace_logs_record = []
            log_msg = await self.client.make_request_log_msg(self.request)
            log_msg_var.set(log_msg)
            logger.info(log_msg, event_name='request')

        # 下一个循环体
        async def _next_send(message: Message) -> None:
            if message.get("type") == "http.response.start":
                response_info.headers = Headers(raw=message.get("headers"))
                response_info.status_code = message.get("status")
            # 解析响应体内容信息
            elif message.get("type") == "http.response.body":
                if body := message.get("body"):
                    response_info.body += body.decode("utf8")
                # print('响应头吧', response_info)
                response = Response(content=response_info.body,
                                    status_code=response_info.status_code,
                                    headers=dict(response_info.headers)
                                    )
                await self.after_request(request=self.request, token=token, response=response)

            await send(message)

        try:
            await self.app(scope, receive, _next_send)
        finally:
            pass
